package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"sort"
	"strconv"
	"strings"
	"sync/atomic"
	"time"

	"github.com/tealeg/xlsx/v3"
)

var (
	// output 指定输出格式（txt/excel）
	output = flag.String("o", "", "Excel输出路径（例如：/path/to/report.xlsx）")

	// procs 指定并发工作协程数
	// 默认值为CPU核心数，可通过 -p 参数调整
	procs = flag.Int("p", runtime.NumCPU(), "并发处理协程数")

	// exts 限制统计的文件扩展名
	// 示例：-ext=go,md,txt 仅统计Go/Markdown/文本文件
	exts = flag.String("ext", "", "文件扩展名白名单（逗号分隔）")
)

func ShouldExamine(ext string) bool {
	if *exts == "" {
		return true
	}
	ext = strings.ToLower("," + strings.TrimPrefix(ext, ".") + ",")
	// 是否已经包含了该扩展名
	return strings.Contains(*exts, ext)
}

// FileWorker 文件处理协程
func FileWorker(files chan string, result chan CountByExt, progress *int64) {
	total := make(CountByExt)
	defer func() { result <- total }()

	for file := range files {
		count, err := CountLines(file)
		if err != nil {
			log.Printf("文件 %s 统计失败: %v", file, err)
			continue
		}
		log.Println("debug:", count)
		total.Add(count)
		atomic.AddInt64(progress, 1)
	}
}

// 遍历目录
func IterateDir(work chan string, roots ...string) {
	walk := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		if info.IsDir() {
			filename := info.Name()
			if len(filename) > 1 && filename[0] == '.' {
				return filepath.SkipDir
			}
			return nil
		}

		// 带~认为临时文件，不计入
		if strings.Contains(filepath.Base(path), "~") {
			return nil
		}

		if !ShouldExamine(filepath.Ext(path)) {
			return nil
		}
		work <- path
		return nil
	}

	for _, root := range roots {
		err := filepath.Walk(root, walk)
		if err != nil {
			log.Println(err)
		}
	}

	close(work)
}

// GenerateExcel 生成Excel统计报表
func GenerateExcel(data CountByExt, filename string) error {
	file := xlsx.NewFile()
	sheet, err := file.AddSheet("代码统计")
	if err != nil {
		return fmt.Errorf("创建工作表失败: %w", err)
	}

	// 设置表头
	headerRow := sheet.AddRow()
	headerRow.AddCell().SetString("文件类型")
	headerRow.AddCell().SetString("文件数量")
	headerRow.AddCell().SetString("代码行数")
	headerRow.AddCell().SetString("空白行数")

	// 填充统计数据
	for ext, stats := range data {
		dataRow := sheet.AddRow()
		dataRow.AddCell().SetString(ext)
		dataRow.AddCell().SetInt(stats.FileCount)
		dataRow.AddCell().SetInt(stats.Code)
		dataRow.AddCell().SetInt(stats.Blank)
	}

	// 确保目录存在
	if err := os.MkdirAll(filepath.Dir(filename), 0755); err != nil {
		return fmt.Errorf("目录创建失败: %w", err)
	}
	if err := file.Save(filename); err != nil {
		return fmt.Errorf("文件保存失败: %w", err)
	}
	return nil
}

func main() {
	flag.Parse()

	paths := flag.Args()
	if len(paths) == 0 {
		paths = []string{"."}
	}

	if *exts != "" {
		*exts = "," + *exts + ","
	}

	var total CountByExt
	var progress int64

	fmtProgress := func() string {
		v := atomic.LoadInt64(&progress)
		return "	files " + strconv.Itoa(int(v))
	}

	work := make(chan string, 100)
	results := make(chan CountByExt, *procs)

	go IterateDir(work, paths...)

	for i := 0; i < *procs; i++ {
		go FileWorker(work, results, &progress)
	}

	for N := *procs; N > 0; {
		select {
		case <-time.After(100 * time.Millisecond):
			line := fmtProgress()
			backspace := strings.Repeat("\r", len(line))
			fmt.Print(line, backspace)
		case result := <-results:
			for _, s := range result {
				log.Println("debug:", s) // 调试用，输出每个文件的统计结果
				if total == nil {
					total = make(CountByExt)
				}
				total.Add(s)
			}
			N--
		}
	}

	fmt.Println(fmtProgress())
	fmt.Println()

	if total == nil {
		total = make(CountByExt)
	}

	counts := make(Counts, 0, len(total))
	for _, c := range total {
		counts = append(counts, c)
	}

	sort.Sort(ByCode{counts})

	summary := Count{}

	fmt.Printf("%-12s %12s %12s %12s %12s\n", "extension", "files", "binary", "blank", "code")
	fmt.Println(strings.Repeat("-", 12*5+4))
	for _, count := range counts {
		summary.Add(count)
		fmt.Printf("%-12s %12d %12d %12d %12d\n", count.Ext, count.FileCount, count.Binary, count.Blank, count.Code)
	}
	fmt.Println(strings.Repeat("-", 12*5+4))
	fmt.Printf("%-12s %12d %12d %12d %12d\n", "summary", summary.FileCount, summary.Binary, summary.Blank, summary.Code)

	// 在程序最后生成Excel文件
	if *output != "" {
		if !strings.HasSuffix(*output, ".xlsx") {
			*output += ".xlsx"
		}
		filename := filepath.Clean(*output)
		if err := GenerateExcel(total, filename); err != nil {
			log.Fatalf("生成Excel失败: %v", err)
		}
		fmt.Printf("\n报表已生成: %s\n", filename)
	}
}
